#!/usr/bin/env python
print "importing libraries"
import sm
import kalibr_common as kc
from kalibr_imu_camera_calibration import *

import numpy as np
import argparse
import signal
import sys   

# make numpy print prettier
np.set_printoptions(suppress=True)

def signal_exit(signal, frame):
    print
    sm.logWarn("Shutting down! (CTRL+C)")
    sys.exit(1)

#helper to constrain certain arguments to be specified only once
class Once(argparse.Action):
    def __call__(self, parser, namespace, values, option_string = None):
        if getattr(namespace, self.dest) is not None:
            msg = '{o} can only be specified once'.format(o = option_string)
            raise argparse.ArgumentError(None, msg)
        setattr(namespace, self.dest, values)

def parseArgs():
    class KalibrArgParser(argparse.ArgumentParser):
        def error(self, message):
            self.print_help()
            sm.logError('%s' % message)
            sys.exit(2)
        def format_help(self):
            formatter = self._get_formatter()
            formatter.add_text(self.description)
            formatter.add_usage(self.usage, self._actions,
                                self._mutually_exclusive_groups)
            for action_group in self._action_groups:
                formatter.start_section(action_group.title)
                formatter.add_text(action_group.description)
                formatter.add_arguments(action_group._group_actions)
                formatter.end_section()
            formatter.add_text(self.epilog)
            return formatter.format_help()     
        
    usage = """
    Example usage to calibrate a camera system against an IMU using an aprilgrid.
    Temporal calibration is enabled by default.
    
    %(prog)s --bag MYROSBAG.bag --cam camchain.yaml --imu imu.yaml \\
             --target aprilgrid.yaml 
    
    camchain.yaml: is the camera-system calibration output of the multiple-camera
                   calibratin tool (kalibr_calibrate_cameras)
    
    example aprilgrid.yaml:       |  example imu.yaml: (ADIS16448)
        target_type: 'aprilgrid'  |      accelerometer_noise_density: 0.006  
        tagCols: 6                |      accelerometer_random_walk: 0.0002
        tagRows: 6                |      gyroscope_noise_density: 0.0004
        tagSize: 0.088            |      gyroscope_random_walk: 4.0e-06
        tagSpacing: 0.3           |      update_rate: 200.0"""    

    #setup the argument list
    parser = KalibrArgParser(description='Calibrate the spatial and temporal parameters of an IMU to a camera chain.', usage=usage)
    
    #data source
    groupData = parser.add_argument_group('Dataset source')
    groupData.add_argument('--bag', dest='bagfile', nargs=1, help='Ros bag file containing image and imu data (rostopics specified in the yamls)', action=Once, required=True)
    groupData.add_argument('--bag-from-to', metavar='bag_from_to', type=float, nargs=2, help='Use the bag data starting from up to this time [s]')
    groupData.add_argument('--perform-synchronization',  action='store_true', dest='perform_synchronization', \
                          help='Perform a clock synchronization according to \'Clock synchronization algorithms for network measurements\' by Zhang et al. (2002).')
    
    #configuration files
    groupCam = parser.add_argument_group('Camera system configuration')
    groupCam.add_argument('--cams', dest='chain_yaml', help='Camera system configuration as yaml file', action=Once)
    groupCam.add_argument('--recompute-camera-chain-extrinsics',  action='store_true', dest='recompute_chain_extrinsics', \
                          help='Recompute the camera chain extrinsics. This option is exclusively recommended for debugging purposes in order to identify problems with the camera chain extrinsics.')
    groupCam.add_argument('--reprojection-sigma', type=float, default=1.0, dest='reprojection_sigma', help='Standard deviation of the distribution of reprojected corner points [px]. (default: %(default)s)', required=False)
    
    groupImu = parser.add_argument_group('IMU configuration')
    groupImu.add_argument('--imu', dest='imu_yamls', nargs='+', help='Yaml files holding the IMU noise parameters. The first IMU will be the reference IMU.', required=True, action=Once)
    groupImu.add_argument('--imu-delay-by-correlation', action='store_true', dest='estimate_imu_delay', \
                          help='Estimate the delay between multiple IMUs by correlation. By default, no temporal calibration between IMUs will be performed.', required=False)      
    groupImu.add_argument('--imu-models', nargs='+', dest='imu_models', help='The IMU models to estimate. Currently supported are \'calibrated\', \'scale-misalignment\' and \'scale-misalignment-size-effect\'.', action=Once)
    
    groupTarget = parser.add_argument_group('Calibration target')
    groupTarget.add_argument('--target', dest='target_yaml', help='Calibration target configuration as yaml file', required=True, action=Once)
    
    #optimization options
    groupOpt = parser.add_argument_group('Optimization options')
    groupOpt.add_argument('--no-time-calibration', action='store_true', dest='no_time', help='Disable the temporal calibration', required=False)      
    groupOpt.add_argument('--max-iter', type=int, default=30, dest='max_iter', help='Max. iterations (default: %(default)s)', required=False)
    groupOpt.add_argument('--recover-covariance', action='store_true', dest='recover_cov', help='Recover the covariance of the design variables.', required=False)
    groupOpt.add_argument('--timeoffset-padding', type=float, default=30.e-3, dest='timeoffset_padding', help='Maximum range in which the timeoffset may change during estimation [s] (default: %(default)s)', required=False)

    #Result options  
    outputSettings = parser.add_argument_group('Output options')
    outputSettings.add_argument('--show-extraction', action='store_true', dest='showextraction', help='Show the calibration target extraction. (disables plots)')
    outputSettings.add_argument('--extraction-stepping', action='store_true', dest='extractionstepping', help='Show each image during calibration target extraction  (disables plots)', required=False)
    outputSettings.add_argument('--verbose', action='store_true', dest='verbose', help='Verbose output (disables plots)')
    outputSettings.add_argument('--dont-show-report', action='store_true', dest='dontShowReport', help='Do not show the report on screen after calibration.')
     
    outputSettings.add_argument('--export-poses', action='store_true', dest='exportPoses', help='Also export the optimized poses.')

    #print help if no argument is specified
    if len(sys.argv)==1:
        parser.print_help()
        sys.exit(2)
    
    #Parser the argument list
    try:
        parsed = parser.parse_args()
    except:
        sys.exit(2)
                 
    if parsed.verbose:
        parsed.showextraction = True             
    
    #there is a with the gtk plot widget, so we cant plot if we have opencv windows open...
    #--> disable the plots in these special situations
    if parsed.showextraction or parsed.extractionstepping or parsed.verbose:
        parsed.dontShowReport = True
    
    return parsed

def main():
    # Parse the arguments
    parsed = parseArgs();
    
    #logging modess
    if parsed.verbose:
        sm.setLoggingLevel(sm.LoggingLevel.Debug)
    else:
        sm.setLoggingLevel(sm.LoggingLevel.Info)
        
    signal.signal(signal.SIGINT, signal_exit)
            
    print "Initializing IMUs:"
    if  len(parsed.imu_yamls) == 1 and not parsed.imu_models:
        #Maintain backwards compatibility with previous interface.
        parsed.imu_models = ['calibrated']
    elif len(parsed.imu_yamls) != len(parsed.imu_models):
        sm.logError("Number of IMU yamls and models has to match.")
        sys.exit(2)
    
    imus = list()
    for (imu_yaml, imu_model) in zip(parsed.imu_yamls, parsed.imu_models):
        imuConfig = kc.ImuParameters(imu_yaml)
        imuConfig.printDetails()
        if imu_model == 'calibrated':
            imus.append( sens.IccImu(imuConfig, parsed, isReferenceImu=(not imus), \
                         estimateTimedelay=parsed.estimate_imu_delay) )
        elif imu_model == 'scale-misalignment':
            imus.append( sens.IccScaledMisalignedImu(imuConfig, parsed, isReferenceImu=(not imus), \
                         estimateTimedelay=parsed.estimate_imu_delay) )
        elif imu_model == 'scale-misalignment-size-effect':
            imus.append( sens.IccScaledMisalignedSizeEffectImu(imuConfig, parsed, isReferenceImu=(not imus), \
                         estimateTimedelay=parsed.estimate_imu_delay) )
        else:
            sm.logError("Model {0} is currently unsupported.".format(imu_model))
            sys.exit(2)
    
    #load calibration target configuration 
    targetConfig = kc.CalibrationTargetParameters(parsed.target_yaml)
       
    print "Initializing calibration target:"
    targetConfig.printDetails()
    
    print "Initializing camera chain:"
    chain = kc.CameraChainParameters(parsed.chain_yaml)      
    chain.printDetails()   
    camChain = sens.IccCameraChain(chain, targetConfig, parsed)

    #create a calibrator instance
    iCal = IccCalibrator()

    #register sensors with calibrator
    iCal.registerCamChain(camChain)
    for imu in imus:
        iCal.registerImu(imu)
        if imu is not imus[0]:
            imu.findOrientationPrior(imus[0])

    print
    print "Building the problem"
    iCal.buildProblem(splineOrder=6, 
                      poseKnotsPerSecond=100, 
                      biasKnotsPerSecond=50, 
                      doPoseMotionError=False,
                      doBiasMotionError=True,
                      blakeZisserCam=-1,
                      huberAccel=-1,
                      huberGyro=-1,
                      noTimeCalibration=parsed.no_time,
                      noChainExtrinsics=(not parsed.recompute_chain_extrinsics),
                      maxIterations=parsed.max_iter,
                      timeOffsetPadding=parsed.timeoffset_padding,
                      verbose = parsed.verbose)

    print
    print "Before Optimization"
    print "==================="
    util.printErrorStatistics(iCal)
    
    print
    print "Optimizing..."
    iCal.optimize(maxIterations=parsed.max_iter, recoverCov=parsed.recover_cov)

    print
    print "After Optimization (Results)"
    print "=================="
    util.printErrorStatistics(iCal)
    util.printResults(iCal, withCov=parsed.recover_cov)

    print
    bagtag = parsed.bagfile[0].translate(None, "<>:/\|?*").replace('.bag', '', 1)
    yamlFilename = "camchain-imucam-" + bagtag + ".yaml"
    iCal.saveCamChainParametersYaml(yamlFilename)
    print "  Saving camera chain calibration to file: {0}".format(yamlFilename)

    print
    yamlFilename = "imu-" + bagtag + ".yaml"
    iCal.saveImuSetParametersYaml(yamlFilename)
    print "  Saving imu calibration to file: {0}".format(yamlFilename)

    resultFileTxt = "results-imucam-" + bagtag + ".txt"
    util.saveResultTxt(iCal, filename=resultFileTxt)
    print "  Detailed results written to file: {0}".format(resultFileTxt)
    
    print "Generating result report..."    
    reportFile = "report-imucam-" + bagtag + ".pdf"
    util.generateReport(iCal, filename=reportFile, showOnScreen=not parsed.dontShowReport)
    print "  Report written to {0}".format(reportFile)
    print

    if parsed.exportPoses:
        print "Exporting poses..."
        posesFile = "poses-imucam-imu0-" + bagtag + ".csv"
        util.exportPoses(iCal, filename=posesFile)
        print "Poses written to {0}".format(posesFile)
        print

if __name__ == "__main__":
    main()
#     try:
#         main()
#     except Exception,e:
#         sm.logError("Exception: {0}".format(e))
#         sys.exit(-1)
        
